import torch
import math
from torch import nn
from torch.nn import functional as F
from main import config

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(config.vocab_size, config.vocab_size)


    def forward(self, x):
        return self.embedding.forward(x)
    
    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            y = self.forward(idx)
            idx_next = torch.multinomial(y, 1)
